#!/usr/bin/env python3
# -*- coding: utf-8 -*-
###
# @author Loganli
# @email admin@attacker.club
# @Date: 2025-05-04
# @Description: AWS EC2实例自动化部署Zabbix Agent脚本
###

import os
import sys
import json
import argparse
import configparser
import boto3
from botocore.exceptions import ClientError
import paramiko
import yaml
import concurrent.futures

# 终端颜色设置
class Colors:
    BLUE = '\033[0;34m'
    GREEN = '\033[0;32m'
    YELLOW = '\033[0;33m'
    RED = '\033[0;31m'
    BOLD = '\033[1m'
    CYAN = '\033[0;36m'
    NC = '\033[0m'  # 恢复默认

# 日志输出函数
def log_info(message):
    print(f"{Colors.BLUE}[信息]{Colors.NC} {message}")

def log_success(message):
    print(f"{Colors.GREEN}[成功]{Colors.NC} {message}")

def log_warning(message):
    print(f"{Colors.YELLOW}[警告]{Colors.NC} {message}")

def log_error(message):
    print(f"{Colors.RED}[错误]{Colors.NC} {message}")

class AWSZabbixDeployer:
    def __init__(self, profile=None, region=None, zabbix_server=None, ssh_key_path=None):
        """初始化AWS Zabbix部署器"""
        self.profile = profile
        self.region = region
        self.zabbix_server = zabbix_server
        self.ssh_key_path = ssh_key_path
        self.session = None
        self.ec2_client = None
        self.ec2_resource = None
        self.instances = []
        self.ansible_inventory = {
            'all': {
                'hosts': {},
                'vars': {
                    'ansible_python_interpreter': '/usr/bin/python3',
                    'zabbix_server': self.zabbix_server
                }
            },
            'zabbix_agents': {
                'hosts': {}
            }
        }
        
        # 初始化AWS连接
        self._init_aws_connection()
        
    def _init_aws_connection(self):
        """初始化AWS连接"""
        try:
            if self.profile:
                log_info(f"使用配置文件 {self.profile} 连接AWS")
                self.session = boto3.Session(profile_name=self.profile, region_name=self.region)
            else:
                log_info(f"使用默认凭证连接AWS区域 {self.region}")
                self.session = boto3.Session(region_name=self.region)
                
            self.ec2_client = self.session.client('ec2')
            self.ec2_resource = self.session.resource('ec2')
            log_success("AWS连接初始化成功")
        except Exception as e:
            log_error(f"AWS连接初始化失败: {str(e)}")
            sys.exit(1)
    
    def get_instances(self, filters=None):
        """获取EC2实例列表"""
        if filters is None:
            filters = []
            
        try:
            # 添加默认过滤器：只获取运行中的实例
            filters.append({'Name': 'instance-state-name', 'Values': ['running']})
            
            response = self.ec2_client.describe_instances(Filters=filters)
            
            for reservation in response['Reservations']:
                for instance in reservation['Instances']:
                    instance_info = {
                        'id': instance['InstanceId'],
                        'type': instance['InstanceType'],
                        'state': instance['State']['Name'],
                        'private_ip': instance.get('PrivateIpAddress', ''),
                        'public_ip': instance.get('PublicIpAddress', ''),
                        'key_name': instance.get('KeyName', ''),
                        'tags': instance.get('Tags', [])
                    }
                    
                    # 获取实例名称
                    instance_name = ''
                    for tag in instance_info['tags']:
                        if tag['Key'] == 'Name':
                            instance_name = tag['Value']
                            break
                    
                    instance_info['name'] = instance_name
                    self.instances.append(instance_info)
            
            log_success(f"成功获取 {len(self.instances)} 个运行中的EC2实例")
            return self.instances
        except Exception as e:
            log_error(f"获取EC2实例失败: {str(e)}")
            return []
    
    def generate_ansible_inventory(self, output_path='inventory.yml'):
        """生成Ansible清单文件"""
        if not self.instances:
            log_warning("没有可用的EC2实例，无法生成Ansible清单")
            return False
        
        try:
            for instance in self.instances:
                host_vars = {
                    'ansible_host': instance['private_ip'] if instance['private_ip'] else instance['public_ip'],
                    'ansible_user': 'ec2-user',  # 默认用户，可根据AMI类型调整
                    'ansible_ssh_private_key_file': self.ssh_key_path,
                    'zabbix_hostname': instance['name'] if instance['name'] else instance['id']
                }
                
                # 根据实例标签调整SSH用户
                if any(tag['Key'] == 'OS' and 'ubuntu' in tag['Value'].lower() for tag in instance['tags']):
                    host_vars['ansible_user'] = 'ubuntu'
                
                # 添加到清单
                host_name = instance['name'] if instance['name'] else instance['id']
                self.ansible_inventory['all']['hosts'][host_name] = host_vars
                self.ansible_inventory['zabbix_agents']['hosts'][host_name] = None
            
            # 写入YAML文件
            with open(output_path, 'w') as f:
                yaml.dump(self.ansible_inventory, f, default_flow_style=False)
            
            log_success(f"Ansible清单文件已生成: {output_path}")
            return True
        except Exception as e:
            log_error(f"生成Ansible清单文件失败: {str(e)}")
            return False
    
    def generate_ansible_playbook(self, output_path='deploy_zabbix.yml'):
        """生成Ansible playbook文件"""
        try:
            playbook = [
                {
                    'name': 'Deploy Zabbix Agent to AWS EC2 instances',
                    'hosts': 'zabbix_agents',
                    'become': True,
                    'tasks': [
                        {
                            'name': 'Download Zabbix Agent installation script',
                            'get_url': {
                                'url': 'https://gitee.com/attacker/deploy-zabbix/raw/master/agent/zabbix-agent-universal.sh',
                                'dest': '/tmp/zabbix-agent-universal.sh',
                                'mode': '0755'
                            }
                        },
                        {
                            'name': 'Install Zabbix Agent',
                            'shell': '/tmp/zabbix-agent-universal.sh {{ zabbix_server }} {{ zabbix_hostname }} agent binary',
                            'args': {
                                'creates': '/etc/zabbix/zabbix_agentd.conf'
                            }
                        },
                        {
                            'name': 'Ensure Zabbix Agent is running',
                            'service': {
                                'name': 'zabbix-agent',
                                'state': 'started',
                                'enabled': 'yes'
                            }
                        }
                    ]
                }
            ]
            
            # 写入YAML文件
            with open(output_path, 'w') as f:
                yaml.dump(playbook, f, default_flow_style=False)
            
            log_success(f"Ansible playbook文件已生成: {output_path}")
            return True
        except Exception as e:
            log_error(f"生成Ansible playbook文件失败: {str(e)}")
            return False
    
    def run_ansible_playbook(self, inventory_path='inventory.yml', playbook_path='deploy_zabbix.yml'):
        """执行Ansible playbook"""
        try:
            cmd = f"ansible-playbook -i {inventory_path} {playbook_path}"
            log_info(f"执行命令: {cmd}")
            os.system(cmd)
            return True
        except Exception as e:
            log_error(f"执行Ansible playbook失败: {str(e)}")
            return False
    
    def deploy_direct(self, instance_ip, instance_name, ssh_key_path=None):
        """直接通过SSH部署Zabbix Agent（不使用Ansible）"""
        if not ssh_key_path and not self.ssh_key_path:
            log_error("未指定SSH密钥路径")
            return False
        
        key_path = ssh_key_path if ssh_key_path else self.ssh_key_path
        
        try:
            # 创建SSH客户端
            ssh = paramiko.SSHClient()
            ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
            
            # 连接到实例
            log_info(f"正在连接到实例 {instance_name} ({instance_ip})...")
            ssh.connect(instance_ip, username='ec2-user', key_filename=key_path)
            
            # 构建安装命令
            install_cmd = f"curl -s https://gitee.com/attacker/deploy-zabbix/raw/master/agent/zabbix-agent-universal.sh | bash -s {self.zabbix_server} {instance_name} agent binary"
            
            # 执行命令
            log_info(f"正在执行Zabbix Agent安装命令...")
            stdin, stdout, stderr = ssh.exec_command(install_cmd)
            
            # 输出结果
            for line in stdout:
                print(line.strip())
            
            # 检查错误
            errors = stderr.read().decode()
            if errors:
                log_warning(f"安装过程中出现警告或错误: {errors}")
            
            # 关闭连接
            ssh.close()
            log_success(f"Zabbix Agent已成功部署到实例 {instance_name}")
            return True
        except Exception as e:
            log_error(f"直接部署Zabbix Agent失败: {str(e)}")
            return False

def setup_aws_config():
    """设置AWS配置"""
    config = configparser.ConfigParser()
    config_path = os.path.expanduser('~/.aws/config')
    
    # 检查配置文件是否存在
    if not os.path.exists(config_path):
        os.makedirs(os.path.dirname(config_path), exist_ok=True)
        config['default'] = {
            'region': 'ap-northeast-1',
            'output': 'json'
        }
        with open(config_path, 'w') as f:
            config.write(f)
        log_info(f"已创建AWS配置文件: {config_path}")
    else:
        log_info(f"已存在AWS配置文件: {config_path}")
    
    # 检查凭证文件是否存在
    credentials_path = os.path.expanduser('~/.aws/credentials')
    if not os.path.exists(credentials_path):
        log_warning("未找到AWS凭证文件，请确保已配置AWS凭证")
        aws_access_key = input("请输入AWS Access Key ID: ")
        aws_secret_key = input("请输入AWS Secret Access Key: ")
        
        config = configparser.ConfigParser()
        config['default'] = {
            'aws_access_key_id': aws_access_key,
            'aws_secret_access_key': aws_secret_key
        }
        
        with open(credentials_path, 'w') as f:
            config.write(f)
        
        log_success(f"已创建AWS凭证文件: {credentials_path}")
    else:
        log_info(f"已存在AWS凭证文件: {credentials_path}")

def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='AWS EC2实例自动化部署Zabbix Agent')
    parser.add_argument('--profile', help='AWS配置文件名称')
    parser.add_argument('--region', default='ap-northeast-1', help='AWS区域')
    parser.add_argument('--zabbix-server', required=True, help='Zabbix服务器IP地址')
    parser.add_argument('--ssh-key', help='SSH密钥路径')
    parser.add_argument('--tag', help='按标签筛选EC2实例 (格式: key=value)')
    parser.add_argument('--inventory', default='inventory.yml', help='Ansible清单文件输出路径')
    parser.add_argument('--playbook', default='deploy_zabbix.yml', help='Ansible playbook文件输出路径')
    parser.add_argument('--no-run', action='store_true', help='仅生成文件，不自动运行Ansible playbook')
    parser.add_argument('--direct', action='store_true', help='直接通过SSH部署（不使用Ansible）')
    parser.add_argument('--setup-aws', action='store_true', help='设置AWS配置')
    parser.add_argument('--max-workers', type=int, default=5, help='并发部署的最大线程数')
    parser.add_argument('--quiet', action='store_true', help='减少输出信息')
    
    args = parser.parse_args()
    
    # 设置AWS配置
    if args.setup_aws:
        setup_aws_config()
    
    # 检查SSH密钥
    ssh_key_path = args.ssh_key
    if not ssh_key_path:
        # 检查~/.ssh/目录下的.pem文件
        ssh_dir = os.path.expanduser('~/.ssh')
        pem_files = [f for f in os.listdir(ssh_dir) if f.endswith('.pem')]
        
        if pem_files:
            log_info(f"在~/.ssh/目录下找到以下SSH密钥文件:")
            for i, key_file in enumerate(pem_files):
                print(f"{i+1}. {key_file}")
            
            choice = input("请选择要使用的SSH密钥文件 (输入数字): ")
            try:
                index = int(choice) - 1
                if 0 <= index < len(pem_files):
                    ssh_key_path = os.path.join(ssh_dir, pem_files[index])
                    log_info(f"已选择SSH密钥: {ssh_key_path}")
                else:
                    log_error("无效的选择")
                    sys.exit(1)
            except ValueError:
                log_error("请输入有效的数字")
                sys.exit(1)
        else:
            log_error("未找到SSH密钥文件，请使用--ssh-key参数指定")
            sys.exit(1)
    
    # 创建部署器
    deployer = AWSZabbixDeployer(
        profile=args.profile,
        region=args.region,
        zabbix_server=args.zabbix_server,
        ssh_key_path=ssh_key_path
    )
    
    # 构建过滤器
    filters = []
    if args.tag:
        try:
            key, value = args.tag.split('=')
            filters.append({'Name': f'tag:{key}', 'Values': [value]})
        except ValueError:
            log_error("标签格式错误，应为key=value")
            sys.exit(1)
    
    # 获取实例
    instances = deployer.get_instances(filters)
    
    if not instances:
        log_error("未找到符合条件的EC2实例")
        sys.exit(1)
    
    # 显示找到的实例
    log_info("找到以下EC2实例:")
    
    # 美化表格输出
    fmt = "{:<3} {:<20} {:<20} {:<15} {:<15}"
    print(f"\n{Colors.BOLD}{fmt.format('ID', '实例名称', '实例ID', 'IP地址', '密钥名称')}{Colors.NC}")
    print("─" * 80)
    
    for i, instance in enumerate(instances):
        ip = instance['private_ip'] if instance['private_ip'] else instance['public_ip']
        name = instance['name'] if instance['name'] else 'N/A'
        print(fmt.format(
            f"{i+1}.", 
            name[:18] + '..' if len(name) > 20 else name, 
            instance['id'], 
            ip or 'N/A', 
            instance['key_name'] or 'N/A'
        ))
    print("─" * 80)
    
    # 根据部署方式执行不同操作
    if args.direct:
        # 直接部署模式
        instance_choice = input(f"\n{Colors.BOLD}请选择要部署的实例 (输入数字，多个实例用逗号分隔，全部输入'all'):{Colors.NC} ")
        
        if instance_choice.lower() == 'all':
            selected_indices = range(len(instances))
        else:
            try:
                selected_indices = [int(idx.strip()) - 1 for idx in instance_choice.split(',')]
            except ValueError:
                log_error("请输入有效的数字")
                sys.exit(1)
        
        # 使用线程池并发部署
        import concurrent.futures
        
        selected_instances = []
        for idx in selected_indices:
            if 0 <= idx < len(instances):
                selected_instances.append(instances[idx])
            else:
                log_warning(f"无效的实例索引: {idx+1}")
        
        if not selected_instances:
            log_error("未选择有效的实例")
            sys.exit(1)
        
        confirm = input(f"\n{Colors.BOLD}是否继续部署? (y/n):{Colors.NC} ").strip().lower()
        if confirm != 'y':
            log_info("用户取消部署")
            sys.exit(0)
        
        success_count = 0
        failed_count = 0
        
        with concurrent.futures.ThreadPoolExecutor(max_workers=args.max_workers) as executor:
            future_to_instance = {}
            
            for instance in selected_instances:
                ip = instance['private_ip'] if instance['private_ip'] else instance['public_ip']
                if not ip:
                    log_warning(f"实例 {instance['name']} ({instance['id']}) 没有可用IP地址，跳过")
                    failed_count += 1
                    continue
                
                future = executor.submit(
                    deployer.deploy_direct, 
                    ip, 
                    instance['name'] or instance['id'],
                    ssh_key_path
                )
                future_to_instance[future] = instance
            
            for future in concurrent.futures.as_completed(future_to_instance):
                instance = future_to_instance[future]
                try:
                    result = future.result()
                    if result:
                        success_count += 1
                    else:
                        failed_count += 1
                except Exception as e:
                    log_error(f"部署异常 {instance['name']} ({instance['id']}): {str(e)}")
                    failed_count += 1
        
        # 显示部署结果
        print(f"\n{Colors.BOLD}{Colors.CYAN}▶ 开始批量部署 Zabbix Agent...{Colors.NC}")
        print("─" * 80)
        print("\n" + "═" * 80)
        log_info(f"部署完成: {Colors.GREEN}成功 {success_count} 个{Colors.NC}, {Colors.RED}失败 {failed_count} 个{Colors.NC}")
        print("═" * 80)
    else:
        # Ansible部署模式 - 默认自动运行
        # 生成Ansible清单文件
        deployer.generate_ansible_inventory(args.inventory)
        
        # 生成Ansible playbook文件
        deployer.generate_ansible_playbook(args.playbook)
        
        # 默认自动运行Ansible playbook，除非指定--no-run
        if not args.no_run:
            print("\n" + "═" * 80)
            log_info(f"{Colors.BOLD}{Colors.CYAN}▶ 正在执行Ansible部署...{Colors.NC}")
            print("═" * 80 + "\n")
            deployer.run_ansible_playbook(args.inventory, args.playbook)
        else:
            log_info(f"Ansible文件已生成，可以使用以下命令运行:")
            print(f"\n    {Colors.BOLD}ansible-playbook -i {args.inventory} {args.playbook}{Colors.NC}\n")

if __name__ == '__main__':
    main()
